#include <stdint.h>
#include "bootloader.h"
#include "ff.h"
#include "flash_port.h"

typedef enum {
    UDISK_UPDATE_IDLE = 0,
    UDISK_UPDATE_CHECK,
    UDISK_UPDATE_START,
    UDISK_UPDATE_IN_PROGRESS,
    UDISK_UPDATE_COMPLETE
} udisk_update_state_t;

#define U_DETECT_TIMEOUT 100000 // 5 seconds timeout for USB detection
#define USB_DRIVE_PATH "1:"
#define APP_BIN_PATH "1:/app.bin"


uint8_t udisk_status = 0;
uint8_t udisk_update_state = 0; 
static uint8_t flash_eraser_flag = 0;
static uint32_t start_time = 0;
static FATFS UsbFS;
static FIL firmware_file;
static FILINFO fno;
extern FRESULT res;

HAL_StatusTypeDef iap_write_appbin(uint32_t flash_addr, uint8_t *file_buffer, uint32_t bytes_read);
void JumpToApp(uint32_t appAddr);


void udisk_update(void)
{

    uint8_t file_buffer[512]={0};
    uint32_t bytes_read;
    unsigned long total_bytes_read = 0;
    uint32_t flash_addr = APP_FLASH_ADDR;

    switch (udisk_update_state)
    {
        case UDISK_UPDATE_IDLE:
            udisk_update_state = UDISK_UPDATE_CHECK;
            start_time = HAL_GetTick();
            /* code */
            break;
        case UDISK_UPDATE_CHECK:
            if (HAL_GetTick() - start_time > U_DETECT_TIMEOUT) {
                udisk_update_state = UDISK_UPDATE_COMPLETE;
            }
            else {
                if(udisk_status == 1)
                {
                    res = f_mount(&UsbFS, USB_DRIVE_PATH, 1);
                    if(res == FR_OK)
                    {
                        res = f_stat(APP_BIN_PATH, &fno);
                        if(res == FR_OK)
                        {
                            udisk_update_state = UDISK_UPDATE_START;
                        }
                        else
                        {
                            // udisk_update_state = UDISK_UPDATE_COMPLETE;
                            f_mount(&UsbFS, USB_DRIVE_PATH, 0);
                        }

                    }
                }
            }
            break;
        case UDISK_UPDATE_START:
            res = f_open(&firmware_file, APP_BIN_PATH, FA_READ);
            if(res == FR_OK) {
                udisk_update_state = UDISK_UPDATE_IN_PROGRESS;
                // printf("Starting firmware update from: %s\n", APP_BIN_PATH);
            }
            else
            {
                udisk_update_state = UDISK_UPDATE_COMPLETE;
                // printf("Failed to open update file: %s\n", APP_BIN_PATH);
            }
            break;
        case UDISK_UPDATE_IN_PROGRESS:
            // Here you would typically read the file and perform the update
            while (total_bytes_read < fno.fsize) {
                // 读取数据块到缓冲区
                res = f_read(&firmware_file, file_buffer, 512, &bytes_read);
                if (res != FR_OK || bytes_read == 0) {
                    // 读取出错或到达文件末尾
                    // printf("Error reading file or end of file reached.\n");
                    break;
                }
                // 写入到FLASH
                // printf("Writing %lu bytes to flash at address 0x%08lX\n", bytes_read, flash_addr);
                iap_write_appbin(flash_addr, file_buffer, bytes_read);

                // 更新计数器和地址
                total_bytes_read += bytes_read;
                flash_addr += bytes_read;
                
                // 显示进度
                // printf("Progress: %lu/%lu bytes (%.2f%%)\n", total_bytes_read, fno.fsize, (total_bytes_read * 100.0) / fno.fsize);
            }
            f_close(&firmware_file);
            udisk_update_state = UDISK_UPDATE_COMPLETE;
            // printf("Firmware update completed successfully.\n");

            break;
        case UDISK_UPDATE_COMPLETE:
            udisk_update_state = UDISK_UPDATE_IDLE;
            // printf("Rebooting to new application...\n");
            JumpToApp(APP_FLASH_ADDR);
            break;
        default:
            break;
    }

}



/**
 * @brief 将数据写入Flash
 * @param flash_addr  目标Flash地址（必须为扇区起始地址）
 * @param file_buffer 源数据缓冲区
 * @param bytes_read  数据长度（字节）
 * @retval HAL_OK 成功，其他值失败
 */
HAL_StatusTypeDef iap_write_appbin(uint32_t flash_addr, uint8_t *file_buffer, uint32_t bytes_read) {

    // 1. 检查地址和长度是否合法
    if (flash_addr < APP_FLASH_ADDR || 
        (flash_addr + bytes_read) > (APP_FLASH_ADDR + APP_FLASH_SIZE)) {
        return HAL_ERROR;
    }

    // 擦除目标扇区（STM32F4需按扇区擦除）
    if(flash_eraser_flag == 0)
    {
        flash_port_erase(APP_FLASH_ADDR, APP_FLASH_SIZE);
        flash_eraser_flag = 1;
    }

    flash_port_write(flash_addr, (uint32_t *)file_buffer, bytes_read);

    return HAL_OK;
}



void JumpToApp(uint32_t appAddr) {
    // 2. 复位外设（关键步骤！）
    HAL_RCC_DeInit();           // 复位时钟配置（可选，APP会重新初始化）
    HAL_DeInit();               // 复位所有HAL外设

    __disable_irq();
    /* 关闭滴答定时器，复位到默认值 */
    SysTick->CTRL = 0;
    SysTick->LOAD = 0;
    SysTick->VAL = 0;

    for (uint8_t i = 0; i < 8; i++) {  // 8 个 ICER 寄存器（共 256 个中断）
            NVIC->ICER[i] = 0xFFFFFFFF;    // 禁用中断
            NVIC->ICPR[i] = 0xFFFFFFFF;    // 清除挂起标志
    }
    __enable_irq();

    // 3. 关闭特定外设（按需添加）
    // HAL_DMA_DeInit(&hdma_usart1_tx);  // 关闭DMA
    // HAL_UART_DeInit(&huart1);         // 关闭UART
    // SysTick->CTRL = 0;                // 禁用SysTick

    // 4. 设置堆栈指针并跳转
    uint32_t *pStack = (uint32_t *)appAddr;
    uint32_t *pReset = (uint32_t *)(appAddr + 4);
    __set_MSP(*pStack);               // 设置主堆栈指针
    ((void (*)(void))(*pReset))();    // 跳转到Reset_Handler
    while(1);
}






















































